#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import time
import logging
import argparse
import importlib.util
import torch

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 首先检查并尝试修复脚本路径
script_dir = os.path.dirname(os.path.abspath(__file__))
if script_dir not in sys.path:
    sys.path.append(script_dir)

# 导入自定义模块
try:
    from experiment_manager import ExperimentManager
    from config import Config
    from utils import get_gpu_memory_map
    from data_processor import DataProcessor
    from model_handler import ModelHandler
except ImportError as e:
    logger.error(f"导入模块失败: {str(e)}")
    logger.info("尝试在当前目录查找模块...")
    # 尝试直接从当前目录导入
    sys.path.insert(0, os.getcwd())
    try:
        from experiment_manager import ExperimentManager
        from config import Config
        from utils import get_gpu_memory_map
        from data_processor import DataProcessor
        from model_handler import ModelHandler
    except ImportError as e2:
        logger.error(f"仍然无法导入模块: {str(e2)}")
        logger.error("请确保所有依赖模块都已正确安装和配置。")
        sys.exit(1)


def parse_arguments():
    """
    解析命令行参数
    
    Returns:
        解析后的参数对象
    """
    parser = argparse.ArgumentParser(description="运行患者对话实验")
    
    # 数据相关参数
    parser.add_argument('--data_dir', type=str, default=None,
                        help='数据目录路径')
    parser.add_argument('--max_samples', type=int, default=None,
                        help='最大样本数量')
    
    # 模型相关参数
    parser.add_argument('--model_path', type=str, default=None,
                        help='模型路径')
    parser.add_argument('--use_cpu', action='store_true',
                        help='使用CPU而不是GPU')
    parser.add_argument('--device', type=str, default=None,
                        help='指定使用的设备，例如: cuda:0, cpu')
    
    # 实验相关参数
    parser.add_argument('--output_dir', type=str, default=None,
                        help='结果输出目录')
    parser.add_argument('--context_windows', type=str, default=None,
                        help='上下文窗口大小，多个值用逗号分隔，例如: 2,4,6,full,semantic')
    parser.add_argument('--batch_size', type=int, default=None,
                        help='批处理大小')
    parser.add_argument('--enable_parallel', action='store_true',
                        help='启用并行处理')
    
    # 其他参数
    parser.add_argument('--config_path', type=str, default=None,
                        help='自定义配置文件路径')
    
    return parser.parse_args()


def setup_environment(args):
    """
    设置实验环境
    
    Args:
        args: 命令行参数
        
    Returns:
        配置对象
    """
    # 加载配置
    if args.config_path and os.path.exists(args.config_path):
        # 尝试加载自定义配置文件
        try:
            spec = importlib.util.spec_from_file_location("custom_config", args.config_path)
            custom_config = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(custom_config)
            config = custom_config.Config()
            logger.info(f"成功加载自定义配置文件: {args.config_path}")
        except Exception as e:
            logger.error(f"加载自定义配置文件失败: {str(e)}")
            logger.info("使用默认配置...")
            config = Config()
    else:
        # 使用默认配置
        config = Config()
    
    # 根据命令行参数更新配置
    if args.data_dir:
        config.DATA_DIR = args.data_dir
    if args.model_path:
        config.MODEL_PATH = args.model_path
    if args.output_dir:
        config.OUTPUT_DIR = args.output_dir
    if args.batch_size:
        config.BATCH_SIZE = args.batch_size
    if args.max_samples:
        config.MAX_SAMPLES = args.max_samples
    
    # # 设置并行处理
    # if args.enable_parallel:
    #     config.ENABLE_PARALLEL_PROCESSING = True
    #     # 可以在这里设置并行工作线程数
    #     if hasattr(args, 'max_workers') and args.max_workers:
    #         config.MAX_PARALLEL_WORKERS = args.max_workers
    
    # 设置设备
    if args.use_cpu:
        config.DEVICE = 'cpu'
    elif args.device:
        config.DEVICE = args.device
    # else:
    #     # 自动选择GPU
    #     config.DEVICE = setup_gpu_device()
    
    # 确保输出目录存在
    if not os.path.exists(config.OUTPUT_DIR):
        try:
            os.makedirs(config.OUTPUT_DIR)
            logger.info(f"创建输出目录: {config.OUTPUT_DIR}")
        except Exception as e:
            logger.error(f"创建输出目录失败: {str(e)}")
            # 使用临时目录作为备选
            import tempfile
            temp_dir = tempfile.mkdtemp()
            logger.warning(f"将使用临时目录作为输出目录: {temp_dir}")
            config.OUTPUT_DIR = temp_dir
    
    return config

def validate_environment(config):
    """
    验证实验环境
    
    Args:
        config: 配置对象
        
    Returns:
        验证是否通过
    """
    # 检查数据目录
    if not os.path.exists(config.DATA_DIR):
        logger.error(f"数据目录不存在: {config.DATA_DIR}")
        return False
    
    # 检查数据文件
    data_files = os.listdir(config.DATA_DIR)
    if not data_files:
        logger.error(f"数据目录为空: {config.DATA_DIR}")
        return False
    
    # 检查模型路径
    if not os.path.exists(config.MODEL_PATH):
        logger.error(f"模型路径不存在: {config.MODEL_PATH}")
        return False
    
    # 打印环境信息
    logger.info(f"实验环境验证通过！")
    logger.info(f"数据目录: {config.DATA_DIR}")
    logger.info(f"模型路径: {config.MODEL_PATH}")
    logger.info(f"输出目录: {config.OUTPUT_DIR}")
    
    return True


def output_experiment_summary(total_samples, success_samples, total_time):
    """
    输出实验摘要
    
    Args:
        total_samples: 总样本数
        success_samples: 成功样本数
        total_time: 总耗时（秒）
    """
    if total_samples == 0:
        logger.warning("没有处理任何样本")
        return
    
    # 计算成功率
    success_rate = (success_samples / total_samples) * 100 if total_samples > 0 else 0
    
    # 计算平均每个样本耗时
    avg_time_per_sample = total_time / total_samples if total_samples > 0 else 0
    
    logger.info("=== 实验摘要 ===")
    logger.info(f"总样本数: {total_samples}")
    logger.info(f"成功样本数: {success_samples}")
    logger.info(f"成功率: {success_rate:.2f}%")
    logger.info(f"总耗时: {total_time:.2f}秒")
    logger.info(f"平均每个样本耗时: {avg_time_per_sample:.4f}秒")
    logger.info("===============")


def main():
    """
    主函数
    """
    # 记录开始时间
    start_time = time.time()
    
    try:
        # 解析命令行参数
        args = parse_arguments()
        
        # 设置实验环境
        config = setup_environment(args)
        
        # 验证实验环境
        if not validate_environment(config):
            logger.error("实验环境验证失败，无法继续")
            sys.exit(1)
        
        # 初始化实验管理器
        experiment_manager = ExperimentManager(config)
        
        # 确定上下文窗口大小
        context_windows = None
        if args.context_windows:
            # 解析上下文窗口大小参数
            context_windows = []
            for window in args.context_windows.split(','):
                window = window.strip()
                if window == 'full' or window == 'semantic':
                    context_windows.append(window)
                else:
                    try:
                        context_windows.append(int(window))
                    except ValueError:
                        logger.warning(f"无效的上下文窗口大小: {window}")
        
        # 执行实验
        logger.info("开始执行实验...")
        experiment_results = experiment_manager.run_experiments(
            context_windows=context_windows,
            num_samples=args.max_samples
        )
        
        # 汇总结果
        total_samples = 0
        success_samples = 0
        
        for result in experiment_results:
            total_samples += result.get('total_samples', 0)
            # 计算成功样本数
            for sample in result.get('results', []):
                if 'error' not in sample:
                    success_samples += 1
        
        # 计算总耗时
        total_time = time.time() - start_time
        
        # 输出实验摘要
        output_experiment_summary(total_samples, success_samples, total_time)
        
        # 实验成功完成
        logger.info("实验已成功完成！")
        return 0
        
    except Exception as e:
        logger.error(f"实验执行过程中发生错误: {str(e)}")
        import traceback
        logger.error(f"详细错误栈: {traceback.format_exc()}")
        return 1


if __name__ == '__main__':
    # 运行主函数并返回退出码
    sys.exit(main())